# -*- coding:utf-8 -*-
"""
作者：谷台阳
日期：2022年10月27日
"""
import cv2
import numpy as np
import skimage
import scipy
import pandas as pd
import time

def rgb2ycbcr(image, only_y=True):
    """
    :param image: rgb
    :param only_y: only for y channel
    :return: result: ycbcr
    """
    image_type = image.dtype
    image.astype(np.float32)
    if image_type != np.uint8:
        image *= 255
    if only_y:
        result = np.dot(image, [0.257, 0.504, 0.098]) + 16
    else:
        result = np.matmul(image, [[0.257, 0.504, 0.098], [-0.148, -0.291, 0.439], [0.439, -0.386, -0.071]]) + [16, 128, 128]
    if image_type == np.uint8:
        result.round()
    else:
        result /= 255
    return result.astype(image_type)

def compare_vifp(ref, dist):
    """
    :param ref: image
    :param dist: image
    :return: vifp: num
    """
    sigma_nsq = 2
    eps = 1e-10

    num = 0.0
    den = 0.0
    for scale in range(1, 5):

        N = 2 ** (4 - scale + 1) + 1
        sd = N / 5.0

        if (scale > 1):
            ref = scipy.ndimage.gaussian_filter(ref, sd)
            dist = scipy.ndimage.gaussian_filter(dist, sd)
            ref = ref[::2, ::2]
            dist = dist[::2, ::2]

        mu1 = scipy.ndimage.gaussian_filter(ref, sd)
        mu2 = scipy.ndimage.gaussian_filter(dist, sd)
        mu1_sq = mu1 * mu1
        mu2_sq = mu2 * mu2
        mu1_mu2 = mu1 * mu2
        sigma1_sq = scipy.ndimage.gaussian_filter(ref * ref, sd) - mu1_sq
        sigma2_sq = scipy.ndimage.gaussian_filter(dist * dist, sd) - mu2_sq
        sigma12 = scipy.ndimage.gaussian_filter(ref * dist, sd) - mu1_mu2

        sigma1_sq[sigma1_sq < 0] = 0
        sigma2_sq[sigma2_sq < 0] = 0

        g = sigma12 / (sigma1_sq + eps)
        sv_sq = sigma2_sq - g * sigma12

        g[sigma1_sq < eps] = 0
        sv_sq[sigma1_sq < eps] = sigma2_sq[sigma1_sq < eps]
        sigma1_sq[sigma1_sq < eps] = 0

        g[sigma2_sq < eps] = 0
        sv_sq[sigma2_sq < eps] = 0

        sv_sq[g < 0] = sigma2_sq[g < 0]
        g[g < 0] = 0
        sv_sq[sv_sq <= eps] = eps

        num += np.sum(np.log10(1 + g * g * sigma1_sq / (sv_sq + sigma_nsq)))
        den += np.sum(np.log10(1 + sigma1_sq / sigma_nsq))

    vifp = num / den

    if np.isnan(vifp):
        return 1.0
    else:
        return vifp


if __name__ == '__main__':

    data = [{"order": []}, {"psnr": []}, {"ssim": []}, {"vif": []}]
    start = time.perf_counter()
    for i in range(1, 9951):
        image_message = cv2.imread("../data/image_message/im{}.jpg".format(10000+i))
        image_original = cv2.imread("../data/image_original/im{}.jpg".format(10000+i))
        image_message = rgb2ycbcr(image_message)
        image_original = rgb2ycbcr(image_original)

        psnr_data = skimage.metrics.peak_signal_noise_ratio(image_original, image_message)
        ssim_data = skimage.metrics.structural_similarity(image_original, image_message)
        vif_data = compare_vifp(image_original, image_message)

        data[0]["order"].append(10000+i)
        data[1]["psnr"].append(ssim_data)
        data[2]["ssim"].append(psnr_data)
        data[3]["vif"].append(vif_data)

        c = int(i / 9950 * 100)
        a = "*" * c
        b = "-" * (100 - c)
        last = time.perf_counter() - start
        print("\r{:^3.0f}%[{}->{}]{:.2f}s".format(c, a, b, last), end="")

    df0 = pd.DataFrame(data[0])
    df1 = pd.DataFrame(data[1])
    df3 = pd.DataFrame(data[2])
    df4 = pd.DataFrame(data[3])

    with pd.ExcelWriter("eval.xlsx") as writer:
        df0.to_excel(writer, sheet_name='eval', index=False, startrow=1, startcol=0, header=False)
        df1.to_excel(writer, sheet_name='eval', index=False, startrow=0, startcol=1, header=True)
        df3.to_excel(writer, sheet_name='eval', index=False, startrow=0, startcol=2, header=True)
        df4.to_excel(writer, sheet_name='eval', index=False, startrow=0, startcol=3, header=True)